import csv
from gensim.models import Word2Vec
from openpyxl.styles import Font, Alignment
from excel_utl import ExcelTool


class ModelTrainer(Word2Vec):
    def __init__(self,
                 sentences,
                 vector_size=128,
                 window=5,
                 min_count=2,
                 workers=8,
                 sg=1,
                 epochs=100,
                 alpha=0.025,
                 min_alpha=0.0001,
                 hs=0,
                 compute_loss=True,
                 callbacks=[]):
        super().__init__(
            sentences=sentences,
            vector_size=vector_size,
            window=window,
            min_count=min_count,
            workers=workers,
            sg=sg,
            epochs=epochs,
            alpha=alpha,
            min_alpha=min_alpha,
            hs=hs,
            compute_loss=compute_loss,
            callbacks=callbacks)

    def expand_and_save(self, core_words, output_file, similarity_threshold=0.5, similar_words_count=80):
        """
        扩展核心词并保存其相似词到文件。
        """
        print("开始扩展核心词并保存结果...")
        similar_words = {}

        for word in core_words:
            if word in self.wv:
                # 获取相似词及其相似度
                raw_similar_words = self.wv.most_similar(word, topn=200)
                # 过滤相似词
                filtered_similar_words = []
                for sim_word, sim_score in raw_similar_words:
                    if sim_score > similarity_threshold and len(filtered_similar_words) <= similar_words_count:
                        filtered_similar_words.append(
                            (sim_word, round(sim_score, 3)))
                # 保存过滤后的结果
                similar_words[word] = filtered_similar_words
                print(
                    f"核心词 '{word}' 的相似词（余弦相似度 > {similarity_threshold}）：{similar_words[word]}")
            else:
                print(f"核心词 '{word}' 不在词汇表中，跳过。")

        # 保存到csv
        with open(output_file, 'w', encoding='utf-8', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(["核心词", "扩展词", "相似度"])
            for word in similar_words:
                for sim_word, sim_score in similar_words[word]:
                    writer.writerow([word, sim_word, sim_score])

        print(f"核心词及扩展词已保存到文件：{output_file}")
        return similar_words

    def calc_deti_and_count_word(self, document_dict: dict, core_word: str, similar_words: dict, company: str, year: str) -> tuple:
        """
        计算数据要素文本指数
        document_dict: 公司词数据字典（公司：年：词数据）
        similar_words: 核心词以及扩展词字典（核心词：(扩展词,相似度)）
        company: (i)公司名称（如平安银行）
        year: (t)年份（如2022）
        """

        # 公司在指定年份的年报词数据列表
        document = document_dict[company][year]

        # 1.计算分母
        sum = 0.
        # 扩展词元组(扩展词,相似度)列表
        expand_word_tuple_list = similar_words[core_word]
        # 扩展词字典(扩展词:相似度)
        expand_word_dict = {
            expand_word: similarity for expand_word, similarity in expand_word_tuple_list}
        # 核心词相似度设为1
        expand_word_dict[core_word] = 1.0
        # 当前核心词+扩展词
        expand_words = set(expand_word_dict.keys())
        expand_words.add(core_word)
        # 统计当前核心词和扩展词在公司年报中出现的次数
        citm_dict = {word: 0 for word in expand_words}
        for word in document:
            if word in expand_words:
                citm_dict[word] += 1
        # 计算每个词出现的次数 * 相似度 再求和
        for word, citm in citm_dict.items():
            simitm = expand_word_dict[word]
            sum += citm * simitm

        # 2.计算分子
        company_word_count = len(document)

        # 3.设置K
        K = 1000

        # 计算最终结果detiit
        detiit = (sum / company_word_count) * K
        return detiit, citm_dict

    def calc_and_save_excel(self, documents_dict: dict, similar_words: dict, output_file: str):
        """
        计算deti,保存excel
        """
        excel_tool = ExcelTool()
        need_change_sheet = False
        for core_word in similar_words.keys():
            # 创建一个sheet
            if need_change_sheet:
                excel_tool.switch_sheet(core_word)
            else:
                excel_tool.set_sheet_title(core_word)

            header = ["公司", "年份", core_word + "_deti", "词频求和"]
            similar_word_list = [word for word,
                                 similar in similar_words[core_word]]
            header.append(core_word)
            header.extend(similar_word_list)

            header_dict_list = []
            for company_name, company_dict in documents_dict.items():
                for year in company_dict.keys():
                    deti, count_dict = self.calc_deti_and_count_word(
                        documents_dict, core_word, similar_words, company_name, year)
                    header_dict = {key: None for key in header}
                    header_dict["公司"] = company_name
                    header_dict["年份"] = year
                    header_dict[core_word + "_deti"] = deti
                    header_dict["词频求和"] = sum(count_dict.values())
                    for word, count in count_dict.items():
                        header_dict[word] = count
                    header_dict_list.append(header_dict)

            content = []
            for header_dict in header_dict_list:
                content.append(list(header_dict.values()))

            excel_tool.write_data(header, content)

            header_font = Font(bold=True)
            alignment = Alignment(horizontal="center", vertical="center")
            header_cell_range = excel_tool.get_cell_range(len(header), 1)
            excel_tool.apply_styles(
                header_cell_range, font=header_font, alignment=alignment)
            content_cell_range = excel_tool.get_cell_range(
                len(header), len(content), start_row=2)
            excel_tool.apply_styles(content_cell_range, alignment=alignment)

            # 自动调整列宽
            excel_tool.auto_adjust_columns_width()

            need_change_sheet = True

        # 保存文件
        excel_tool.save(output_file)
        print(f"Excel文件已保存到：{output_file}")

        # 关闭工作簿
        excel_tool.close()
